package com.ds.infrastructure.hyperspace.container.flow.handler.loadbalance;

import com.ds.infrastructure.hyperspace.common.Container;
import com.ds.infrastructure.hyperspace.common.EndpointNode;
import com.ds.infrastructure.hyperspace.common.config.TargetNodeConfig;
import com.ds.infrastructure.hyperspace.common.flow.DuplexFlowContext;
import com.ds.infrastructure.hyperspace.common.flow.FlowHandlerPipeline;
import com.ds.infrastructure.hyperspace.common.flow.TargetNodeContext;
import com.ds.infrastructure.hyperspace.common.flow.TargetNodeWrapper;
import com.ds.infrastructure.hyperspace.common.flow.handler.LoadBalancePreFlowHandler;
import com.ds.infrastructure.hyperspace.container.flow.context.HyperspaceDuplexFlowContext;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.CollectionUtils;

import java.util.List;
import java.util.Map;

/**
 * @author  Weishaoying
 */
@Slf4j
public class RoundRobinWithWeightPreHandler implements LoadBalancePreFlowHandler<HyperspaceDuplexFlowContext> {

    @Override
    public EndpointNode<TargetNodeConfig> select(HyperspaceDuplexFlowContext duplexFlowContext,
            List<EndpointNode<TargetNodeConfig>> nodeList,
            Map<String, ? extends TargetNodeWrapper> nodeContextMap) {
        if (CollectionUtils.isEmpty(nodeList)) {
            return null;
        }
        if (nodeList.size() == 1) {
            return nodeList.get(0);
        }

        int totalWeight = 0;
        EndpointNode<TargetNodeConfig> bestConfigEndpointNode = null;
        TargetNodeContext bestConfigTargetNodeContext = null;
        EndpointNode<TargetNodeConfig> bestEndpointNode = null;
        TargetNodeContext bestTargetNodeContext = null;
        for (EndpointNode<TargetNodeConfig> currentNode : nodeList) {
            TargetNodeWrapper currentTargetNodeWrapper = nodeContextMap.get(currentNode.getKey());
            if (currentTargetNodeWrapper == null) {
                log.warn("TargetNodeContext not found, round robin with weight ignore!");
                continue;
            }

            TargetNodeContext currentTargetNodeContext = currentTargetNodeWrapper.getTargetNodeContext();
            if (bestConfigTargetNodeContext == null ||
                    currentTargetNodeContext.getConfigWeight() >
                            bestConfigTargetNodeContext.getConfigWeight()) {
                bestConfigEndpointNode = currentNode;
                bestConfigTargetNodeContext = currentTargetNodeContext;
            }

            if (!currentNode.isAvailable()) {
                int tempEffectiveWeight = currentTargetNodeContext.getEffectiveWeight() -
                        currentTargetNodeContext.getConfigWeight();
                if (tempEffectiveWeight < 0) {
                    tempEffectiveWeight = 0;
                }
                currentTargetNodeContext.setEffectiveWeight(tempEffectiveWeight);
                continue;
            }

            int effectiveWeight = currentTargetNodeContext.getEffectiveWeight();
            int currentWeight = currentTargetNodeContext.getCurrentWeight() + effectiveWeight;
            totalWeight += effectiveWeight;

            currentTargetNodeContext.setCurrentWeight(currentWeight);
            if (effectiveWeight < currentTargetNodeContext.getConfigWeight()) {
                currentTargetNodeContext.setEffectiveWeight(++effectiveWeight);
            } else if (effectiveWeight > currentTargetNodeContext.getConfigWeight()) {
                currentTargetNodeContext.setEffectiveWeight(currentTargetNodeContext.getConfigWeight());
            }
            if (bestEndpointNode == null || currentWeight >
                    nodeContextMap.get(bestEndpointNode.getKey())
                            .getTargetNodeContext().getCurrentWeight()) {
                bestEndpointNode = currentNode;
                bestTargetNodeContext = currentTargetNodeContext;
            }
        }

        if (bestEndpointNode == null) {
            return bestConfigEndpointNode;
        }
        bestTargetNodeContext.setCurrentWeight(bestTargetNodeContext.getCurrentWeight() - totalWeight);
        return bestEndpointNode;
    }

    @Override
    public String handlerName() {
        return "roundRobinWithWeightPreHandler";
    }

    @Override
    public void initHandlerParameters(Container container, FlowHandlerPipeline pipeline, String[] handlerParams) {
        // no need implements
    }
    @Override
    public void handleInput(DuplexFlowContext duplexFlowContext, Object data) {
        // no need implements
    }
}
